package com.biubiux2.handler;

import com.biubiux2.message.MessageType;
import com.biubiux2.security.DESUtil;
import com.biubiux2.security.RSAUtil;
import io.netty.buffer.ByteBuf;
import io.netty.buffer.PooledByteBufAllocator;
import io.netty.channel.ChannelFutureListener;
import io.netty.channel.ChannelHandlerContext;
import io.netty.handler.codec.MessageToMessageCodec;
import io.netty.handler.codec.MessageToMessageDecoder;
import io.netty.handler.codec.MessageToMessageEncoder;
import lombok.extern.slf4j.Slf4j;

import javax.crypto.Cipher;
import javax.crypto.SecretKey;
import java.nio.charset.StandardCharsets;
import java.security.PublicKey;
import java.util.LinkedList;
import java.util.List;

/**
 * @author noexcs
 * @since 8/10/2022 7:38 PM
 */
@Slf4j
public class ServerEncryptHandler extends MessageToMessageCodec<ByteBuf, ByteBuf> {

    private Cipher encryptCipher;

    private Cipher decryptCipher;

    private boolean negotiated = false;

    private final LinkedList<ByteBuf> msgWhilePrenegotiated = new LinkedList<>();

    /**
     * @see MessageToMessageEncoder#write(io.netty.channel.ChannelHandlerContext, java.lang.Object, io.netty.channel.ChannelPromise) 会进行release
     */
    @Override
    protected void encode(ChannelHandlerContext ctx, ByteBuf msg, List<Object> out) throws Exception {
        if (negotiated) {
            byte[] bytes = new byte[msg.readableBytes()];
            msg.readBytes(bytes);
            byte[] encrypted = encryptCipher.doFinal(bytes);

            ByteBuf encryptedBytes = PooledByteBufAllocator.DEFAULT.buffer(encrypted.length);
            encryptedBytes.writeBytes(encrypted);
            out.add(encryptedBytes);
        } else {
            msgWhilePrenegotiated.addLast(msg.retainedDuplicate());
        }
    }

    /**
     * @see MessageToMessageDecoder#channelRead(io.netty.channel.ChannelHandlerContext, java.lang.Object) 会进行release
     */
    @Override
    protected void decode(ChannelHandlerContext ctx, ByteBuf msg, List<Object> out) throws Exception {
        if (negotiated) {
            byte[] bytes = new byte[msg.readableBytes()];
            msg.readBytes(bytes);
            byte[] decrypted = decryptCipher.doFinal(bytes);

            ByteBuf decryptedBytes = PooledByteBufAllocator.DEFAULT.buffer(decrypted.length);
            decryptedBytes.writeBytes(decrypted);
            out.add(decryptedBytes);
        } else {
            // 密钥未协商阶段
            int type = msg.readInt();
            if (type == MessageType.CLIENT_PUBLIC_KEY) {
                // 服务端收到 公钥
                byte[] bytes = new byte[msg.readableBytes()];
                msg.readBytes(bytes);
                PublicKey publicKey = RSAUtil.getPublicKey(new String(bytes, StandardCharsets.UTF_8));

                // 生成对称密钥
                SecretKey secretKey = DESUtil.generateSecretKey();
                encryptCipher = DESUtil.getEncryptCipher(secretKey);
                decryptCipher = DESUtil.getDecryptCipher(secretKey);

                // 加密对称密钥
                String keyStr = DESUtil.convertKeyToString(secretKey);
                String encrypt = RSAUtil.encrypt(keyStr, publicKey);

                byte[] outBytes = encrypt.getBytes(StandardCharsets.UTF_8);
                ByteBuf buffer = PooledByteBufAllocator.DEFAULT.buffer(4 + outBytes.length);
                buffer.writeInt(7);
                buffer.writeBytes(outBytes);
                ctx.writeAndFlush(buffer).addListener((ChannelFutureListener) future -> {
                    // 标记协商成功
                    negotiated = true;
                });
            } else {
                log.debug("密钥未协商");
            }
        }
    }
}
